package cn.jolyoulu.common.core.algorithms.flowcontrol;

import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @Author LuZhouJin
 * @Date 2023/6/1
 * 漏桶限流算法
 * 思路：
 * 当前有一个桶，桶中有一个洞会根据一定的速率流水桶中的水，
 * 在需要处理任务时前往桶中添加相当比例的水量
 * 假如桶溢出，代表已经满负荷运作无法接受更多的任务了
 * 假如桶未溢出，代表还有空余位置处理任务
 */
public class LeakyBucket {
    /**
     * 桶的容量
     */
    private final int capacity;
    /**
     * 桶中漏水速度(单位水量)，即每秒桶中会漏多少水量
     */
    private final int rate;
    /**
     * 当前桶的水量，当水量达到capacity表示桶已满
     */
    private final AtomicInteger water;
    /**
     * 上次计算水量的时间
     * (当前时间 - 上一次时间) * rate = 这次water需要减少多少水量
     */
    private final AtomicReference<Instant> lastLeakTime = new AtomicReference<>();

    public LeakyBucket(int capacity, int rate) {
        this.capacity = capacity;
        this.rate = rate;
        this.water = new AtomicInteger(0);
        this.lastLeakTime.set(Instant.now());
        ;
    }

    /**
     * 尝试往桶中添加 1 水量
     *
     * @return true成功，false失败
     */
    public boolean acquire() {
        return acquire(1);
    }

    /**
     * 尝试往桶中添加指定水量
     *
     * @param tokens 水量
     * @return true成功，false失败
     */
    public boolean acquire(int tokens) {
        //更新水量
        leakWater();
        //原子性地增加水量，并获取当前水量
        int currentWaterLevel = water.addAndGet(tokens);
        //判断添加后水量是否超出capacity
        if (currentWaterLevel <= capacity) {
            return true;
        } else {
            //水量超过容量，回滚水量的增加
            water.getAndAdd(-tokens);
            return false;
        }
    }

    private void leakWater() {
        //获取当前时间
        Instant now = Instant.now();
        Instant oldTime;
        long seconds;
        // CAS循环：确保lastLeakTime的更新是“基于最新值”的
        do {
            oldTime = lastLeakTime.get(); // 读取当前时间
            seconds = now.getEpochSecond() - oldTime.getEpochSecond();
            if (seconds <= 0) return; // 无漏水，直接退出
            // 计算漏水量和新水量
            int leakAmount = (int) (seconds * rate);
            int oldWater = water.get();
            int newWater = Math.max(0, oldWater - leakAmount);
            // 先更新水量（原子操作），再尝试更新时间（CAS）
            if (water.compareAndSet(oldWater, newWater)) {
                // 水量更新成功后，尝试更新时间（即使失败也没关系，下次会重新计算）
                lastLeakTime.compareAndSet(oldTime, now);
                return;
            }
        // 如果水量更新失败（被其他线程改了），重新循环
        } while (true);
    }


    public static void main(String[] args) throws InterruptedException {
        LeakyBucket bucket = new LeakyBucket(1, 1);
        List<Thread> threads = new ArrayList<>();
        for (int i = 0; i < 2; i++) {
            threads.add(new Thread(() -> {
                while (true) {
                    String name = Thread.currentThread().getName();
                    System.out.println(name + "=> 尝试往桶里加水，即申请令牌");
                    if (bucket.acquire(1)) {
                        System.out.println(name + "=> 申请成功");
                    } else {
                        System.out.println(name + "=> 申请失败");
                    }
                    try {
                        TimeUnit.MILLISECONDS.sleep(500); //休眠1秒
                    } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                    }
                }
            }, "Thread-" + i));
        }
        threads.forEach(Thread::start);
        TimeUnit.SECONDS.sleep(60);
    }
}
